#!/usr/bin/env python3
"""
Population Population PBPK Modeling

This script generates realistic population PK curves using population-derived
parameters, following the four-mode pattern from population_pbpk_modeling.py.

Model combinations:
- first_order + minimal_pbpk
- first_order + full_pbpk
- adam + minimal_pbpk
- adam + full_pbpk

Author: BlackCat@CPPO
Version: 1.0.0
"""

import numpy as np
import matplotlib.pyplot as plt
import sys
import os
import json
from datetime import datetime

# Add src path to system path
project_root = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, project_root)

from demo.generate_pbpk_parameters import PopulationParameterGenerator
from demo.population_pbpk_modeling import PopulationPBPKSimulator


class PopulationPopulationPBPK:
    """Population PBPK modeling with population-derived parameters"""
    
    def __init__(self, output_dir='../output/population_population'):
        self.output_dir = output_dir
        self.parameter_generator = PopulationParameterGenerator()
        self.ensure_output_dir()
    
    def ensure_output_dir(self):
        """Create output directory if it doesn't exist"""
        os.makedirs(self.output_dir, exist_ok=True)
    
    def run_all_models(self, n_subjects=100):
        """
        Run all four model combinations for population modeling
        
        Args:
            n_subjects: Number of subjects in population
        """
        print("=" * 80)
        print("Population PBPK Modeling - All Four Modes")
        print("=" * 80)
        print(f"Population size: {n_subjects}")
        
        # Define all four model combinations
        model_combinations = [
            ('first_order', 'minimal_pbpk'),
            ('first_order', 'full_pbpk'),
            ('adam', 'minimal_pbpk'),
            ('adam', 'full_pbpk')
        ]
        
        all_results = {}
        
        for abs_model, dist_model in model_combinations:
            try:
                print(f"\nRunning {abs_model} + {dist_model}...")
                
                # Get population parameters
                population_params = self.parameter_generator.get_population_parameters(
                    n_subjects=n_subjects
                )
                
                # Run population simulation
                population_results = self.run_population_simulation(
                    abs_model, dist_model, population_params
                )
                
                key = f"{abs_model}_{dist_model}"
                all_results[key] = population_results
                
                # Save results
                self.save_population_results(key, population_results)
                
                print(f"  Completed {len(population_results)} subjects")
                
            except Exception as e:
                print(f"  Error: {str(e)}")
                continue
        
        # Generate plots
        self.plot_population_curves(all_results)
        self.plot_population_comparisons(all_results)
        
        return all_results
    
    def run_population_simulation(self, abs_model, dist_model, population_params):
        """Run population simulation for a specific model combination"""
        simulator = PopulationPBPKSimulator(n_subjects=len(population_params))
        
        results = []
        
        for i, params in enumerate(population_params):
            try:
                # Prepare ADAM parameters if needed
                if abs_model == 'adam':
                    params.update({
                        'use_adam': True,
                        'solubility': 1.0,
                        'permeability': 1e-4,
                        'particle_size': 50.0,
                        'density': 1.2,
                        'molecular_weight': 300.0,
                        'pKa': 7.4,
                        'logP': 2.0
                    })
                
                # Run individual simulation
                result = simulator.run_subject_simulation(
                    subject_id=params['subject_id'],
                    absorption_model=abs_model,
                    distribution_model=dist_model,
                    params=params
                )
                
                results.append(result)
                
                if (i + 1) % 20 == 0:
                    print(f"    Progress: {i + 1}/{len(population_params)} subjects")
                
            except Exception as e:
                print(f"    Subject {params['subject_id']} failed: {str(e)}")
                continue
        
        return results
    
    def save_population_results(self, model_key, results):
        """Save population results for a specific model"""
        filename = os.path.join(self.output_dir, f'population_{model_key}.json')
        
        # Extract key data for JSON
        time_points = results[0]['time_points'].tolist() if results else []
        
        # Calculate population statistics
        if results:
            all_concentrations = np.array([r['plasma_concentration'] for r in results])
            
            population_data = {
                'model_key': model_key,
                'n_subjects': len(results),
                'time_points': time_points,
                'median_concentration': np.median(all_concentrations, axis=0).tolist(),
                'p5_concentration': np.percentile(all_concentrations, 5, axis=0).tolist(),
                'p25_concentration': np.percentile(all_concentrations, 25, axis=0).tolist(),
                'p75_concentration': np.percentile(all_concentrations, 75, axis=0).tolist(),
                'p95_concentration': np.percentile(all_concentrations, 95, axis=0).tolist(),
                'pk_parameters': {
                    'Cmax': [r['pk_metrics']['C_max'] for r in results],
                    'Tmax': [r['pk_metrics']['T_max'] for r in results],
                    'AUC': [r['pk_metrics']['AUC'] for r in results],
                    't_half': [r['pk_metrics']['t_half'] for r in results]
                }
            }
            
            # Calculate summary statistics
            for param, values in population_data['pk_parameters'].items():
                population_data['pk_parameters'][param] = {
                    'mean': float(np.mean(values)),
                    'std': float(np.std(values)),
                    'median': float(np.median(values)),
                    'min': float(np.min(values)),
                    'max': float(np.max(values)),
                    'cv': float(np.std(values) / np.mean(values) * 100) if np.mean(values) > 0 else 0
                }
        else:
            population_data = {'model_key': model_key, 'n_subjects': 0}
        
        with open(filename, 'w') as f:
            json.dump(population_data, f, indent=2)
    
    def plot_population_curves(self, all_results):
        """Plot population PK curves for all model combinations"""
        model_combinations = [
            ('first_order', 'minimal_pbpk'),
            ('first_order', 'full_pbpk'),
            ('adam', 'minimal_pbpk'),
            ('adam', 'full_pbpk')
        ]
        
        fig, axes = plt.subplots(2, 2, figsize=(16, 12))
        axes = axes.flatten()
        
        colors = {
            'first_order_minimal_pbpk': 'blue',
            'first_order_full_pbpk': 'red',
            'adam_minimal_pbpk': 'green',
            'adam_full_pbpk': 'purple'
        }
        
        for idx, (abs_model, dist_model) in enumerate(model_combinations):
            ax = axes[idx]
            model_key = f"{abs_model}_{dist_model}"
            
            if model_key in all_results and all_results[model_key]:
                results = all_results[model_key]
                
                # Extract data
                time_points = results[0]['time_points']
                all_concentrations = np.array([r['plasma_concentration'] for r in results])
                
                # Plot median and percentiles
                median_conc = np.median(all_concentrations, axis=0)
                p5 = np.percentile(all_concentrations, 5, axis=0)
                p25 = np.percentile(all_concentrations, 25, axis=0)
                p75 = np.percentile(all_concentrations, 75, axis=0)
                p95 = np.percentile(all_concentrations, 95, axis=0)
                
                ax.plot(time_points, median_conc, 'b-', linewidth=3, label='Median')
                ax.fill_between(time_points, p5, p95, alpha=0.3, color='blue', label='5th-95th percentile')
                ax.fill_between(time_points, p25, p75, alpha=0.5, color='blue', label='25th-75th percentile')
                
                # Add statistics text
                Cmax_values = [r['pk_metrics']['C_max'] for r in results]
                AUC_values = [r['pk_metrics']['AUC'] for r in results]
                
                stats_text = f'N = {len(results)}\n'
                stats_text += f'Median Cmax = {np.median(Cmax_values):.2f} mg/L\n'
                stats_text += f'Median AUC = {np.median(AUC_values):.1f} mg·h/L\n'
                stats_text += f'CV(Cmax) = {np.std(Cmax_values)/np.mean(Cmax_values)*100:.1f}%'
                
                ax.text(0.02, 0.98, stats_text, transform=ax.transAxes,
                       verticalalignment='top', bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))
            
            ax.set_xlabel('Time (h)')
            ax.set_ylabel('Plasma Concentration (mg/L)')
            ax.set_title(f'{abs_model.replace("_", " ").title()} + {dist_model.replace("_", " ").title()}')
            ax.grid(True, alpha=0.3)
        
        plt.tight_layout()
        filename = os.path.join(self.output_dir, 'population_pk_curves.png')
        plt.savefig(filename, dpi=300, bbox_inches='tight')
        plt.close()
        print(f"Saved population curves plot: {filename}")
    
    def plot_population_comparisons(self, all_results):
        """Plot comparisons between different model combinations"""
        model_combinations = [
            ('first_order', 'minimal_pbpk'),
            ('first_order', 'full_pbpk'),
            ('adam', 'minimal_pbpk'),
            ('adam', 'full_pbpk')
        ]
        
        fig, axes = plt.subplots(2, 2, figsize=(15, 12))
        
        # Prepare data for boxplots
        labels = []
        Cmax_data = []
        AUC_data = []
        Tmax_data = []
        
        for abs_model, dist_model in model_combinations:
            model_key = f"{abs_model}_{dist_model}"
            if model_key in all_results and all_results[model_key]:
                results = all_results[model_key]
                labels.append(f"{abs_model}\n{dist_model}")
                
                Cmax_data.append([r['pk_metrics']['C_max'] for r in results])
                AUC_data.append([r['pk_metrics']['AUC'] for r in results])
                Tmax_data.append([r['pk_metrics']['T_max'] for r in results])
        
        # Plot PK parameter comparisons
        ax1 = axes[0, 0]
        if Cmax_data:
            ax1.boxplot(Cmax_data, tick_labels=labels)
            ax1.set_ylabel('Cmax (mg/L)')
            ax1.set_title('Cmax Distribution by Model')
            ax1.grid(True, alpha=0.3)
        
        ax2 = axes[0, 1]
        if AUC_data:
            ax2.boxplot(AUC_data, tick_labels=labels)
            ax2.set_ylabel('AUC (mg·h/L)')
            ax2.set_title('AUC Distribution by Model')
            ax2.grid(True, alpha=0.3)
        
        ax3 = axes[1, 0]
        if Tmax_data:
            ax3.boxplot(Tmax_data, tick_labels=labels)
            ax3.set_ylabel('Tmax (h)')
            ax3.set_title('Tmax Distribution by Model')
            ax3.grid(True, alpha=0.3)
        
        # Plot median curves comparison
        ax4 = axes[1, 1]
        colors = plt.cm.Set1(np.linspace(0, 1, len(model_combinations)))
        
        for idx, (abs_model, dist_model) in enumerate(model_combinations):
            model_key = f"{abs_model}_{dist_model}"
            if model_key in all_results and all_results[model_key]:
                results = all_results[model_key]
                time_points = results[0]['time_points']
                
                all_concentrations = np.array([r['plasma_concentration'] for r in results])
                median_conc = np.median(all_concentrations, axis=0)
                
                ax4.plot(time_points, median_conc, color=colors[idx], linewidth=2,
                        label=model_key.replace('_', ' + '))
        
        ax4.set_xlabel('Time (h)')
        ax4.set_ylabel('Median Concentration (mg/L)')
        ax4.set_title('Population Median Curves Comparison')
        ax4.grid(True, alpha=0.3)
        ax4.legend()
        
        plt.tight_layout()
        filename = os.path.join(self.output_dir, 'population_comparisons.png')
        plt.savefig(filename, dpi=300, bbox_inches='tight')
        plt.close()
        print(f"Saved population comparisons plot: {filename}")
    
    def generate_summary_report(self, all_results):
        """Generate comprehensive summary report"""
        report_file = os.path.join(self.output_dir, 'population_summary_report.txt')
        
        model_combinations = [
            ('first_order', 'minimal_pbpk'),
            ('first_order', 'full_pbpk'),
            ('adam', 'minimal_pbpk'),
            ('adam', 'full_pbpk')
        ]
        
        with open(report_file, 'w') as f:
            f.write("Population PBPK Modeling Summary Report\n")
            f.write("=" * 50 + "\n")
            f.write(f"Generated on: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
            f.write(f"Population size: {max([len(r) for r in all_results.values()]) if all_results else 0}\n")
            f.write(f"Output directory: {self.output_dir}\n\n")
            
            for abs_model, dist_model in model_combinations:
                model_key = f"{abs_model}_{dist_model}"
                f.write(f"\n{model_key.replace('_', ' + ').title()}:\n")
                f.write("-" * 40 + "\n")
                
                if model_key in all_results and all_results[model_key]:
                    results = all_results[model_key]
                    
                    Cmax_values = [r['pk_metrics']['C_max'] for r in results]
                    AUC_values = [r['pk_metrics']['AUC'] for r in results]
                    Tmax_values = [r['pk_metrics']['T_max'] for r in results]
                    
                    f.write(f"Subjects: {len(results)}\n")
                    f.write(f"Cmax (mg/L): {np.mean(Cmax_values):.3f} ± {np.std(Cmax_values):.3f}\n")
                    f.write(f"AUC (mg·h/L): {np.mean(AUC_values):.1f} ± {np.std(AUC_values):.1f}\n")
                    f.write(f"Tmax (h): {np.mean(Tmax_values):.1f} ± {np.std(Tmax_values):.1f}\n")
                    f.write(f"CV(Cmax): {np.std(Cmax_values)/np.mean(Cmax_values)*100:.1f}%\n")
                    f.write(f"CV(AUC): {np.std(AUC_values)/np.mean(AUC_values)*100:.1f}%\n")
                    f.write(f"Range Cmax: {np.min(Cmax_values):.3f} - {np.max(Cmax_values):.3f} mg/L\n")
                    f.write(f"Range AUC: {np.min(AUC_values):.1f} - {np.max(AUC_values):.1f} mg·h/L\n")
                else:
                    f.write("No valid results\n")
        
        print(f"Summary report saved: {report_file}")


def main():
    """Main function"""
    print("Population PBPK Modeling")
    print("=" * 50)
    
    # Create runner
    runner = PopulationPopulationPBPK()
    
    # Run all four model combinations
    results = runner.run_all_models(n_subjects=100)
    
    # Generate summary report
    runner.generate_summary_report(results)
    
    print("\n" + "=" * 50)
    print("Modeling completed!")
    print(f"Results saved in: {runner.output_dir}")
    print("=" * 50)


if __name__ == "__main__":
    main()